from flask import Blueprint, request, jsonify, current_app, send_file
from flask_jwt_extended import jwt_required, get_jwt_identity
from models import User, Message, MediaFile, Friendship, db
from utils.encryption import EncryptionUtils
from flask import current_app
from extensions import socketio
import uuid
import datetime
from datetime import timezone
from sqlalchemy import and_, or_
from typing import List, Dict, Any
import json
import os
import io
from werkzeug.utils import secure_filename
import logging
import base64

logger = logging.getLogger(__name__)

messages_bp = Blueprint('messages', __name__, url_prefix='/api/messages')

def get_simple_master_key(user_id: str) -> bytes:
    """简化版主密钥获取（用于演示环境）"""
    import hashlib
    key = hashlib.sha256(f"demo_key_{user_id}".encode()).digest()[:32]
    return key

def get_user_by_identifier(identifier: str) -> User:
    """通过用户名或UUID获取用户"""
    try:
        user_uuid = uuid.UUID(identifier)
        return User.query.get(user_uuid)
    except ValueError:
        return User.query.filter_by(username=identifier).first()

@messages_bp.route('/contacts', methods=['GET'])
@jwt_required()
def get_contacts():
    """获取用户的好友列表"""
    try:
        current_user_id = get_jwt_identity()
        current_user_uuid = uuid.UUID(current_user_id)
        
        # 只获取用户的好友
        friends = Friendship.get_friends(current_user_uuid)
        
        contacts_list = []
        for friend in friends:
            contacts_list.append({
                'id': str(friend.id),
                'username': friend.username,
                'online': friend.online,
                'last_active': friend.last_login.isoformat() if friend.last_login else None
            })
        
        return jsonify(contacts_list), 200
        
    except Exception as e:
        logger.error(f"Get contacts error: {str(e)}", exc_info=True)
        return jsonify({"error": "Failed to get contacts"}), 500

@messages_bp.route('/<string:receiver_identifier>', methods=['GET'])
@jwt_required()
def get_messages(receiver_identifier: str):
    """获取对话消息"""
    try:
        sender_id = get_jwt_identity()
        if not receiver_identifier:
            return jsonify({"error": "Receiver identifier is required"}), 400
        
        receiver = get_user_by_identifier(receiver_identifier)
        if not receiver:
            return jsonify({"error": "Receiver not found"}), 404
        
        sender_uuid = uuid.UUID(sender_id)
        receiver_uuid = receiver.id
        
        # 获取after参数用于增量同步
        after_timestamp = request.args.get('after')
        
        # 构建基础查询
        query = db.session.query(Message).filter(
            or_(
                and_(Message.sender_id == sender_uuid, Message.receiver_id == receiver_uuid),
                and_(Message.sender_id == receiver_uuid, Message.receiver_id == sender_uuid)
            )
        )
        
        # 如果提供了after参数，只获取该时间之后的消息
        if after_timestamp:
            try:
                after_dt = datetime.datetime.fromisoformat(after_timestamp.replace('Z', '+00:00'))
                query = query.filter(Message.timestamp > after_dt)
            except ValueError:
                return jsonify({"error": "Invalid after timestamp format"}), 400
        
        # 按时间升序排序
        messages = query.order_by(Message.timestamp.asc()).all()
        
        if not messages:
            return jsonify([]), 200
        
        # 简化解密处理 - 直接返回明文
        decrypted_messages = []
        for msg in messages:
            try:
                # 尝试从二进制内容解码为字符串
                if msg.encrypted_content:
                    try:
                        if isinstance(msg.encrypted_content, (bytes, bytearray)):
                            content = msg.encrypted_content.decode('utf-8', errors='replace')
                        elif isinstance(msg.encrypted_content, str):
                            content = msg.encrypted_content
                        else:
                            content = str(msg.encrypted_content)
                    except (UnicodeDecodeError, AttributeError):
                        content = str(msg.encrypted_content)
                else:
                    content = "[空消息]"
                
                decrypted_messages.append({
                    'id': str(msg.id),
                    'sender_id': str(msg.sender_id),
                    'receiver_id': str(msg.receiver_id),
                    'message_type': msg.message_type,
                    'content': content,
                    'timestamp': msg.timestamp.isoformat(),
                    'file_ext': msg.file_ext,
                    'is_read': msg.is_read
                })
            except Exception as e:
                decrypted_messages.append({
                    'id': str(msg.id),
                    'sender_id': str(msg.sender_id),
                    'receiver_id': str(msg.receiver_id),
                    'message_type': msg.message_type,
                    'content': f'[处理错误: {str(e)}]',
                    'timestamp': msg.timestamp.isoformat(),
                    'file_ext': msg.file_ext,
                    'is_read': msg.is_read
                })
        
        return jsonify(decrypted_messages), 200
        
    except Exception as e:
        logger.error(f"Get messages error: {str(e)}", exc_info=True)
        return jsonify({"error": "Failed to get messages"}), 500

@messages_bp.route('/', methods=['POST'])
@jwt_required()
def send_message():
    """发送消息API（支持用户名和UUID）"""
    try:
        sender_id = get_jwt_identity()
        data = request.get_json()
        
        receiver_identifier = data.get('receiver_id') or data.get('receiver_username')
        
        # 处理前端发送的消息格式
        message_data = data.get('message', {})
        if isinstance(message_data, dict):
            message_content = message_data.get('content', '')
            message_type = message_data.get('type', 'text')
        else:
            message_content = data.get('content', '') or str(message_data)
            message_type = data.get('type', 'text')
        
        if not receiver_identifier:
            return jsonify({"error": "Receiver ID or username is required"}), 400
            
        if not message_content:
            return jsonify({"error": "Message content is required"}), 400
        
        # 获取接收者
        receiver = get_user_by_identifier(receiver_identifier)
        if not receiver or not receiver.is_active:
            return jsonify({"error": "Receiver not found or inactive"}), 404
        
        sender_uuid = uuid.UUID(sender_id)
        receiver_uuid = receiver.id
        
        # 将字符串内容编码为二进制存储
        message = Message(
            sender_id=sender_uuid,
            receiver_id=receiver_uuid,
            message_type=message_type,
            encrypted_content=message_content.encode('utf-8'),  # 编码为二进制
            message_key=b'',  # 二进制字段
            nonce='',  # 字符串字段
            tag='',    # 字符串字段
            file_ext=None
        )
        
        db.session.add(message)
        db.session.commit()
        
        # 发送WebSocket消息
        try:
            socket_message = {
                'message_id': str(message.id),
                'sender_id': str(sender_uuid),
                'receiver_id': str(receiver_uuid),
                'type': message_type,
                'content': str(message_content),
                'timestamp': message.timestamp.isoformat(),
                'file_ext': None
            }
            
            socketio.emit('new_message', socket_message, room=str(receiver_uuid))
            logger.info(f"Message sent to {receiver_uuid}")
            
        except Exception as ws_error:
            logger.warning(f"WebSocket notification failed: {ws_error}")
        
        return jsonify({
            "status": "success",
            "message_id": str(message.id),
            "timestamp": message.timestamp.isoformat()
        }), 201
        
    except ValueError as e:
        return jsonify({"error": f"Invalid UUID format: {str(e)}"}), 400
    except Exception as e:
        db.session.rollback()
        logger.error(f"Send message error: {str(e)}", exc_info=True)
        return jsonify({"error": f"Message sending failed: {str(e)}"}), 500

@messages_bp.route('/<string:message_id>/delete', methods=['DELETE'])
@jwt_required()
def delete_message(message_id: str):
    """删除消息API（支持撤回）"""
    try:
        user_id = get_jwt_identity()
        message_uuid = uuid.UUID(message_id)
        
        # 查询消息并验证权限
        message = Message.query.get(message_uuid)
        if not message:
            return jsonify({"error": "Message not found"}), 404
        
        # 仅发送者可删除消息（或管理员）
        if str(message.sender_id) != user_id and not get_current_user().role == "admin":
            return jsonify({"error": "Permission denied"}), 403
        
        # 标记消息为已删除（软删除）
        message.deleted = True
        message.deleted_at = datetime.datetime.now(timezone.utc)
        db.session.commit()
        
        # 通知接收者消息已删除
        socketio.emit('message_deleted', {
            "message_id": str(message_uuid),
            "deleted_by": user_id
        }, room=str(message.receiver_id))
        
        return jsonify({"status": "success", "message": "Message deleted"}), 200
        
    except ValueError:
        return jsonify({"error": "Invalid message ID"}), 400
    except Exception as e:
        logger.error(f"Delete message error: {str(e)}", exc_info=True)
        return jsonify({"error": "Failed to delete message"}), 500

@messages_bp.route('/unread-count', methods=['GET'])
@jwt_required()
def get_unread_message_count():
    """获取未读消息计数"""
    try:
        user_id = get_jwt_identity()
        user_uuid = uuid.UUID(user_id)
        
        # 使用子查询优化性能
        unread_count = db.session.query(
            db.func.count(Message.id)
        ).filter(
            and_(
                Message.receiver_id == user_uuid,
                Message.is_read.is_(False),
                Message.deleted.is_(False)
            )
        ).scalar()
        
        return jsonify({"unread_count": unread_count}), 200
        
    except Exception as e:
        logger.error(f"Get unread count error: {str(e)}", exc_info=True)
        return jsonify({"error": "Failed to get unread count"}), 500

@messages_bp.route('/upload-media', methods=['POST'])
@jwt_required()
def upload_media():
    """上传多媒体文件API（支持加密存储）"""
    try:
        user_id = get_jwt_identity()
        
        if 'file' not in request.files:
            return jsonify({'error': 'No file provided'}), 400
            
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': 'Empty filename'}), 400

        # 文件大小检查
        file.seek(0, 2)
        file_size = file.tell()
        file.seek(0)
        
        # 根据文件类型设置大小限制
        max_size_map = {
            'image': 5 * 1024 * 1024,  # 5MB
            'voice': 2 * 1024 * 1024,  # 2MB
            'file': 10 * 1024 * 1024   # 10MB
        }
        
        file_type = 'image' if file.content_type.startswith('image/') else \
                   'voice' if file.content_type.startswith('audio/') else 'file'
        
        max_size = max_size_map.get(file_type, 10 * 1024 * 1024)
        if file_size > max_size:
            return jsonify({'error': f'File too large. Max {max_size//1024//1024}MB'}), 413

        # 生成安全文件名和路径
        from werkzeug.utils import secure_filename
        filename = secure_filename(file.filename)
        file_ext = filename.rsplit('.', 1)[1].lower() if '.' in filename else ''
        
        # 生成唯一文件名
        unique_filename = f"{uuid.uuid4().hex}.{file_ext}"
        upload_dir = os.path.join(current_app.config.get('UPLOAD_FOLDER', 'uploads'), user_id)
        os.makedirs(upload_dir, exist_ok=True)
        file_path = os.path.join(upload_dir, unique_filename)

        # 读取文件内容并加密
        file_content = file.read()
        master_key = get_simple_master_key(user_id)
        
        # 为每个文件生成独立的加密密钥
        file_key = EncryptionUtils.generate_key()
        
        encrypted_content, nonce, tag = EncryptionUtils.aes_encrypt(file_key, file_content)
        
        # 保存加密后的文件
        with open(file_path, 'wb') as f:
            f.write(encrypted_content)
        
        # 保存文件元数据到数据库
        media_file = MediaFile(
            id=uuid.uuid4(),
            user_id=uuid.UUID(user_id),
            filename=filename,
            file_path=file_path,
            file_key=EncryptionUtils.binary_to_hex(file_key),
            nonce=EncryptionUtils.binary_to_hex(nonce),
            tag=EncryptionUtils.binary_to_hex(tag),
            content_type=file.content_type,
            size=file_size
        )
        db.session.add(media_file)
        db.session.commit()

        # 返回加密的文件信息
        return jsonify({
            'status': 'success',
            'media_id': str(media_file.id),
            'filename': filename,
            'content_type': file.content_type,
            'size': file_size,
            'file_key': EncryptionUtils.binary_to_hex(file_key),
            'nonce': EncryptionUtils.binary_to_hex(nonce),
            'tag': EncryptionUtils.binary_to_hex(tag)
        }), 201
        
    except Exception as e:
        db.session.rollback()
        logger.error(f"Media upload failed: {str(e)}")
        return jsonify({'error': 'Media upload failed'}), 500

@messages_bp.route('/media/<string:media_id>', methods=['GET'])
def get_media(media_id: str):
    """获取加密媒体文件"""
    try:
        # 从URL参数或Authorization头获取token
        token = request.args.get('token') or request.headers.get('Authorization', '').replace('Bearer ', '')
        if not token:
            return jsonify({'error': 'Token required'}), 401
            
        try:
            from flask_jwt_extended import decode_token
            decoded_token = decode_token(token)
            user_id = decoded_token['sub']
        except Exception as e:
            return jsonify({'error': 'Invalid token'}), 401
            
        media_uuid = uuid.UUID(media_id)
        
        # 查询媒体文件
        media_file = MediaFile.query.get(media_uuid)
        if not media_file or getattr(media_file, 'is_deleted', False):
            return jsonify({'error': 'Media not found'}), 404
            
        # 验证权限（上传者或相关聊天参与者可访问）
        if str(media_file.user_id) != user_id:
            # 检查用户是否为相关聊天参与者 - 放宽权限检查
            # 只要用户与上传者有聊天记录，就可以访问媒体文件
            user_uuid = uuid.UUID(user_id)
            media_user_uuid = media_file.user_id
            
            # 检查是否有互相发送消息的历史
            related_conversation = Message.query.filter(
                (
                    ((Message.sender_id == user_uuid) & (Message.receiver_id == media_user_uuid)) |
                    ((Message.sender_id == media_user_uuid) & (Message.receiver_id == user_uuid))
                )
            ).first()
            
            if not related_conversation:
                return jsonify({'error': 'Permission denied'}), 403
        
        # 读取加密文件
        if not os.path.exists(media_file.file_path):
            return jsonify({'error': 'File not found on disk'}), 404
            
        with open(media_file.file_path, 'rb') as f:
            encrypted_content = f.read()
        
        # 解密文件
        file_key = EncryptionUtils.hex_to_binary(media_file.file_key)
        nonce = EncryptionUtils.hex_to_binary(media_file.nonce)
        tag = EncryptionUtils.hex_to_binary(media_file.tag)
        
        decrypted_content = EncryptionUtils.aes_decrypt(file_key, encrypted_content, nonce, tag)
        
        # 返回文件内容
        import io
        from flask import send_file
        return send_file(
            io.BytesIO(decrypted_content),
            mimetype=media_file.content_type,
            as_attachment=False,
            download_name=media_file.filename
        )
        
    except Exception as e:
        logger.error(f"Get media failed: {str(e)}")
        return jsonify({'error': 'Failed to get media'}), 500

@messages_bp.route('/download/<media_id>', methods=['GET'])
@jwt_required()
def download_media(media_id):
    """解密并下载媒体文件"""
    current_user_id = get_current_user_id()
    try:
        media_id_uuid = uuid.UUID(media_id)
        media = MediaFile.query.get(media_id_uuid)
        if not media:
            return jsonify({'error': 'Media not found'}), 404

        # 验证文件所有者
        if str(media.user_id) != current_user_id:
            # 检查是否为执法部门授权访问
            if not has_law_enforcement_access():
                return jsonify({'error': 'Permission denied'}), 403

        # 解密文件
        file_key = EncryptionUtils.hex_to_binary(media.file_key)
        nonce = EncryptionUtils.hex_to_binary(media.nonce)
        tag = EncryptionUtils.hex_to_binary(media.tag)
        decrypted_path = media.file_path[:-4]  # 移除.enc后缀

        try:
            EncryptionUtils.file_decrypt(media.file_path, decrypted_path, file_key)
            # 读取文件内容并删除临时文件
            with open(decrypted_path, 'rb') as f:
                file_data = f.read()
            os.remove(decrypted_path)

            return send_file(
                io.BytesIO(file_data),
                mimetype=media.content_type,
                as_attachment=True,
                download_name=media.filename
            )
        except Exception as e:
            logger.error(f"File decryption failed: {str(e)}")
            return jsonify({'error': 'File decryption failed'}), 500
    except ValueError:
        return jsonify({'error': 'Invalid media ID format'}), 400
    except Exception as e:
        logger.error(f"File download failed: {str(e)}")
        return jsonify({'error': 'File download failed'}), 500